package bch.rocketmq.test;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.AbstractQueuedSynchronizer;

/**
 * 基于AQS实现的自定义信号量工具类
 * 支持公平和非公平模式，提供许可证的获取和释放功能
 */
public class CustomSemaphore {

    /**
     * 继承AQS实现同步逻辑
     */
    private abstract static class Sync extends AbstractQueuedSynchronizer {
        private static final long serialVersionUID = 1L;

        Sync(int permits) {
            setState(permits);
        }

        final int getPermits() {
            return getState();
        }

        final int nonfairTryAcquireShared(int acquires) {
            for (; ; ) {
                int available = getState();
                int remaining = available - acquires;
                if (remaining < 0 || compareAndSetState(available, remaining)) {
                    return remaining;
                }
            }
        }

        @Override
        protected final boolean tryReleaseShared(int releases) {
            for (; ; ) {
                int current = getState();
                int next = current + releases;
                if (next < current) // overflow
                {
                    throw new Error("Maximum permit count exceeded");
                }
                if (compareAndSetState(current, next)) {
                    return true;
                }
            }
        }

        final void reducePermits(int reductions) {
            for (; ; ) {
                int current = getState();
                int next = current - reductions;
                if (next > current) // underflow
                {
                    throw new Error("Permit count underflow");
                }
                if (compareAndSetState(current, next)) {
                    return;
                }
            }
        }

        final int drainPermits() {
            for (; ; ) {
                int current = getState();
                if (current == 0 || compareAndSetState(current, 0)) {
                    return current;
                }
            }
        }
    }

    /**
     * 非公平模式同步器
     */
    static final class NonfairSync extends Sync {
        private static final long serialVersionUID = 1L;

        NonfairSync(int permits) {
            super(permits);
        }

        @Override
        protected int tryAcquireShared(int acquires) {
            return nonfairTryAcquireShared(acquires);
        }
    }

    /**
     * 公平模式同步器
     */
    static final class FairSync extends Sync {
        private static final long serialVersionUID = 1L;

        FairSync(int permits) {
            super(permits);
        }

        @Override
        protected int tryAcquireShared(int acquires) {
            for (; ; ) {
                if (hasQueuedPredecessors()) {
                    return -1;
                }
                int available = getState();
                int remaining = available - acquires;
                if (remaining < 0 || compareAndSetState(available, remaining)) {
                    return remaining;
                }
            }
        }
    }

    /**
     * 同步器实例
     */
    private final Sync sync;

    /**
     * 创建具有给定许可数量和非公平设置的信号量
     */
    public CustomSemaphore(int permits) {
        sync = new NonfairSync(permits);
    }

    /**
     * 创建具有给定许可数量和公平设置的信号量
     */
    public CustomSemaphore(int permits, boolean fair) {
        sync = fair ? new FairSync(permits) : new NonfairSync(permits);
    }

    /**
     * 获取一个许可证，如果没有可用许可证则阻塞
     */
    public void acquire() throws InterruptedException {
        sync.acquireSharedInterruptibly(1);
    }

    /**
     * 获取指定数量的许可证
     */
    public void acquire(int permits) throws InterruptedException {
        if (permits < 0) throw new IllegalArgumentException();
        sync.acquireSharedInterruptibly(permits);
    }

    /**
     * 尝试获取一个许可证，不阻塞
     */
    public boolean tryAcquire() {
        return sync.nonfairTryAcquireShared(1) >= 0;
    }

    /**
     * 尝试获取指定数量的许可证，不阻塞
     */
    public boolean tryAcquire(int permits) {
        if (permits < 0) throw new IllegalArgumentException();
        return sync.nonfairTryAcquireShared(permits) >= 0;
    }

    /**
     * 在指定时间内尝试获取一个许可证
     */
    public boolean tryAcquire(long timeout, TimeUnit unit) throws InterruptedException {
        return sync.tryAcquireSharedNanos(1, unit.toNanos(timeout));
    }

    /**
     * 在指定时间内尝试获取指定数量的许可证
     */
    public boolean tryAcquire(int permits, long timeout, TimeUnit unit)
            throws InterruptedException {
        if (permits < 0) throw new IllegalArgumentException();
        return sync.tryAcquireSharedNanos(permits, unit.toNanos(timeout));
    }

    /**
     * 释放一个许可证
     */
    public void release() {
        sync.releaseShared(1);
    }

    /**
     * 释放指定数量的许可证
     */
    public void release(int permits) {
        if (permits < 0) throw new IllegalArgumentException();
        sync.releaseShared(permits);
    }

    /**
     * 获取当前可用许可证数量
     */
    public int availablePermits() {
        return sync.getPermits();
    }

    /**
     * 获取并返回所有立即可用的许可证
     */
    public int drainPermits() {
        return sync.drainPermits();
    }

    /**
     * 减少指定数量的许可证
     */
    protected void reducePermits(int reduction) {
        if (reduction < 0) throw new IllegalArgumentException();
        sync.reducePermits(reduction);
    }

    /**
     * 查询是否有线程等待获取许可证
     */
    public final boolean hasQueuedThreads() {
        return sync.hasQueuedThreads();
    }

    /**
     * 获取等待获取许可证的线程数量的估计值
     */
    public final int getQueueLength() {
        return sync.getQueueLength();
    }

    /**
     * 返回此信号量的字符串表示形式
     */
    public String toString() {
        return super.toString() + "[Permits = " + sync.getPermits() + "]";
    }
}

/**
 * 使用示例和测试类
 */
class CustomSemaphoreDemo {

    public static void main(String[] args) throws InterruptedException {
        // 创建一个有3个许可证的信号量
        CustomSemaphore semaphore = new CustomSemaphore(3);

        System.out.println("初始许可证数量: " + semaphore.availablePermits());

        // 创建并启动多个线程来测试信号量
        for (int i = 0; i < 5; i++) {
            final int threadId = i;
            new Thread(() -> {
                try {
                    System.out.println("线程 " + threadId + " 尝试获取许可证...");
                    semaphore.acquire();
                    System.out.println("线程 " + threadId + " 获得许可证，剩余: " +
                            semaphore.availablePermits());

                    // 模拟工作
                    Thread.sleep(2000);

                    System.out.println("线程 " + threadId + " 释放许可证");
                    semaphore.release();

                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            }).start();
        }

        // 等待一段时间观察结果
        Thread.sleep(1000);
        System.out.println("等待线程数量: " + semaphore.getQueueLength());

        // 测试超时获取
        Thread timeoutThread = new Thread(() -> {
            try {
                System.out.println("超时测试线程尝试在1秒内获取许可证...");
                boolean acquired = semaphore.tryAcquire(1, TimeUnit.SECONDS);
                System.out.println("超时获取结果: " + acquired);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        });

        timeoutThread.start();
        timeoutThread.join();

        Thread.sleep(8000); // 等待所有线程完成
        System.out.println("最终许可证数量: " + semaphore.availablePermits());
    }
}